package com.king.spark.rdd.transform

import org.apache.spark.rdd.RDD
import org.apache.spark.util.AccumulatorV2
import org.apache.spark.{SparkConf, SparkContext}

import scala.collection.mutable

/**
 * @Author wdl
 * @Date 2022/11/22 16:34
 */
object Spark_RDD_Accumulator {

  def main(args: Array[String]): Unit = {

    val sparkConf = new SparkConf().setMaster("local[*]").setAppName("WordCount")

    val sc: SparkContext = new SparkContext(sparkConf)

    val words: RDD[String] = sc.makeRDD( List ("a","a","b","c","b","a") ,2 )
    //使用自定义累加器的步骤:
    //1. 创建累加器
    val wordCountAcc = new WordCountAccumulator

    //2. 注册累加器
    sc.register(wordCountAcc, "wcAcc")

    //3. 使用累加器
    words.foreach(
      word => wordCountAcc.add(word)
    )

    //4. 获取累加器的结果

    val result: mutable.Map[String, Long] = wordCountAcc.value
    println(result)

    sc.stop()
  }

  def main1(args: Array[String]): Unit = {

    val sparkConf = new SparkConf().setMaster("local[*]").setAppName("WordCount")

    val sc: SparkContext = new SparkContext(sparkConf)

    val rdd = sc.makeRDD(List(1,2,3,4,5))
    // 声明累加器
    var sum = sc.longAccumulator("sum");
    rdd.foreach(
      num => {
        // 使用累加器
        sum.add(num)
      }
    )
    // 获取累加器的值
    println("sum = " + sum.value)

    sc.stop()
  }
}

class WordCountAccumulator extends AccumulatorV2[String, mutable.Map[String, Long]]{

  var map : mutable.Map[String, Long] = mutable.Map()

  // 累加器是否为初始状态
  override def isZero: Boolean = map.isEmpty

  // 复制累加器
  override def copy(): AccumulatorV2[String, mutable.Map[String, Long]] = new WordCountAccumulator

  // 重置累加器
  override def reset(): Unit = map.clear()

  // 向累加器中增加数据 (In)
  override def add(v: String): Unit = {
    map.getOrElse(v, 0L) + 1L
    map.put(v, map.getOrElse(v, 0L) + 1L)
  }

  // 合并累加器
  override def merge(other: AccumulatorV2[String, mutable.Map[String, Long]]): Unit = {
    val value1: mutable.Map[String, Long] = other.value
    val map1: mutable.Map[String, Long] = map
    map = map1.foldLeft(value1){
      (a, b) => {
        a.put(b._1, a.getOrElse(b._1, 0L) + b._2)
        a
      }
    }
  }

  override def value: mutable.Map[String, Long] = map
}